import cv2
import numpy as np
import matplotlib.pyplot as plt
from skimage import feature, exposure, transform, morphology
from scipy.signal import find_peaks
from skimage.morphology import footprint_rectangle


def show_img(img, title="Image", cmap='gray'):
    """显示图像"""
    plt.imshow(img, cmap=cmap)
    plt.title(title)
    plt.axis('off')
    plt.show()


def measure_stripe_spacing(image_path, scale_bar_pixels, scale_bar_actual, manual_mode=False):
    """测量电镜图中条纹间隔的主函数，支持自动和手动模式"""
    # 1. 读取图像并利用高斯模糊去噪
    img = cv2.imread(image_path, 0)
    if img is None:
        print(f"错误:无法加载图像，请检查路径:{image_path}")
        return
    blurred = cv2.GaussianBlur(img, (3, 3), 0)

    # 2. 自适应对比度增强
    p2, p98 = np.percentile(blurred, (4, 98))
    enhanced = exposure.rescale_intensity(blurred, in_range=(p2, p98))

    # 3. 自适应边缘检测
    sobel_x = cv2.Sobel(enhanced, cv2.CV_64F, 1, 0, ksize=3)
    sobel_y = cv2.Sobel(enhanced, cv2.CV_64F, 0, 1, ksize=3)
    gradient_mag = np.sqrt(sobel_x ** 2 + sobel_y ** 2)

    grad_mean = np.mean(gradient_mag)
    if grad_mean > 233:
        sigma = 5
    else:
        sigma = 3

    edges = feature.canny(enhanced, sigma=sigma)

    # 4. 显示中间结果
    plt.figure(figsize=(12, 8))
    plt.subplot(221), plt.imshow(img, cmap='gray'), plt.title('原图')
    plt.subplot(222), plt.imshow(enhanced, cmap='gray'), plt.title('对比度增强')
    plt.subplot(223), plt.imshow(edges, cmap='gray'), plt.title('边缘检测结果')
    plt.tight_layout()
    plt.show()

    # 5. 手动框选条纹区域
    print("请在弹出窗口中框选包含条纹的矩形区域（尽量包含多条完整条纹），按回车确认:")
    roi = cv2.selectROI("框选条纹区域", img)
    x, y, w, h = roi
    roi_img = img[y:y + h, x:x + w]
    roi_edges = edges[y:y + h, x:x + w]
    cv2.destroyAllWindows()

    # 准备边缘图像用于后续处理
    roi_edges_uint8 = roi_edges.astype(np.uint8) * 255

    # 6. 检测线条方向
    rotation_angle = 0
    angle_fourier = 0
    angle_hough = 0

    if not manual_mode:
        print("正在自动检测条纹方向...")

        # 霍夫变换检测直线
        lines = cv2.HoughLinesP(
            roi_edges_uint8,
            rho=1,
            theta=np.pi / 180,
            threshold=10,
            minLineLength=min(w, h) * 0.1,  # 动态最小长度
            maxLineGap=3
        )

        # 绘制检测结果
        if lines is not None:
            print(f"检测到 {len(lines)} 条直线")
            img_copied = roi_img.copy()
            for line in lines:
                x1, y1, x2, y2 = line[0]
                cv2.line(img_copied, (x1, y1), (x2, y2), (0, 255, 0), 1)
            show_img(img_copied, "霍夫直线检测结果")

        # 傅里叶变换分析方向
        f = np.fft.fft2(roi_img)
        fshift = np.fft.fftshift(f)
        magnitude_spectrum = 20 * np.log(np.abs(fshift) + 1)  # 修正括号错误

        # 频谱预处理
        magnitude_spectrum = cv2.GaussianBlur(magnitude_spectrum, (5, 5), 0)

        # 创建环形掩码
        h_spec, w_spec = magnitude_spectrum.shape
        center = (w_spec // 2, h_spec // 2)
        y, x = np.indices((h_spec, w_spec))
        r = np.sqrt((x - center[0]) ** 2 + (y - center[1]) ** 2)

        # 动态半径计算
        min_dim = min(h_spec, w_spec)
        inner_radius = max(5, min_dim // 20)  # 自适应内半径
        outer_radius = min_dim // 4  # 自适应外半径
        r_mask = (r > inner_radius) & (r < outer_radius)

        # 计算角度
        angles = np.arctan2(y - center[1], x - center[0]) * 180 / np.pi
        angles_0_180 = np.mod(angles, 180.0)
        angles_0_180[angles_0_180 >= 179.5] -= 180  # 处理边界值

        weights = magnitude_spectrum[r_mask].flatten()
        hist, bins = np.histogram(angles_0_180[r_mask], bins=180, range=(0, 180), weights=weights)

        # 精化峰值搜索
        peak_idx = np.argmax(hist)
        search_range = 10  # 搜索±10个bin
        start = max(0, peak_idx - search_range)
        end = min(len(hist), peak_idx + search_range)
        angle_fourier = bins[start:end][np.argmax(hist[start:end])]

        print(f"傅里叶变换检测方向: {angle_fourier:.2f} 度")

        # 霍夫变换分析方向
        if lines is not None and len(lines) > 3:
            angles_hough = []
            for line in lines:
                x1, y1, x2, y2 = line[0]
                length = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
                if length > min(w, h) * 0.15:  # 只考虑足够长的线
                    angle = np.degrees(np.arctan2(y2 - y1, x2 - x1)) % 180
                    angles_hough.append(angle)

            if angles_hough:
                angle_hough = np.median(angles_hough)
                print(f"霍夫变换检测方向: {angle_hough:.2f} 度")

                # 检查角度一致性
                angle_diff = abs(angle_hough - angle_fourier)
                if angle_diff > 10 and angle_diff < 170:
                    print(f"警告：两种方法角度差异较大 ({angle_diff:.1f}°)，结果可能不可靠")
                    rotation_angle = angle_fourier
                    method = "傅里叶变换（霍夫结果不一致）"
                else:
                    rotation_angle = angle_hough
                    method = "霍夫变换"
            else:
                rotation_angle = angle_fourier
                method = "傅里叶变换（无有效直线）"
        else:
            rotation_angle = angle_fourier
            method = "傅里叶变换（无有效直线）"

        print(f"自动检测结果:条纹方向 {rotation_angle:.2f} 度 ({method})")

        # 可视化角度分布
        plt.figure(figsize=(10, 5))
        plt.bar(bins[:-1], hist, width=0.5, align='edge', alpha=0.7)
        plt.axvline(x=angle_fourier, color='b', linestyle='-', label=f'傅里叶变换 ({angle_fourier:.1f}°)')
        if angle_hough != 0:
            plt.axvline(x=angle_hough, color='g', linestyle='-', label=f'霍夫变换 ({angle_hough:.1f}°)')
        plt.title('角度分布直方图')
        plt.xlabel('角度 (度)')
        plt.ylabel('加权频数')
        plt.legend()
        plt.grid(True)
        plt.show()

        # 手动验证
        manual_check = input("条纹是否水平？需要手动调整角度吗？(y/n):")
        if manual_check.lower() == 'y':
            rotation_angle = float(input("请输入正确角度（度）:"))
    else:
        rotation_angle = float(input("请输入条纹的倾斜角度（度）:"))

    # 7. 旋转图像 (使用OpenCV保持图像完整)
    print(f"旋转图像: {rotation_angle:.2f} 度")
    M = cv2.getRotationMatrix2D((w / 2, h / 2), rotation_angle, 1.0)
    rotated_roi = cv2.warpAffine(roi_img, M, (w, h), flags=cv2.INTER_CUBIC)
    rotated_edges = cv2.warpAffine(roi_edges_uint8, M, (w, h), flags=cv2.INTER_CUBIC)

    # 可视化旋转结果
    plt.figure(figsize=(12, 6))
    plt.subplot(121), plt.imshow(roi_img, cmap='gray'), plt.title('原始ROI')
    plt.subplot(122), plt.imshow(rotated_roi, cmap='gray'), plt.title('旋转后ROI')
    plt.tight_layout()
    plt.show()

    # 8. 投影分析
    print("请选择投影方向:1) 垂直 2) 水平 (默认垂直)")
    choice = input("输入数字(1/2):")
    if choice == '2':
        projection_axis = 0
        projection_dir = "水平"
    else:
        projection_axis = 1
        projection_dir = "垂直"

    print(f"选择{projection_dir}投影分析条纹间隔")
    projection = np.sum(rotated_edges, axis=projection_axis)

    # 优化投影信号
    smoothed_projection = np.convolve(projection, np.ones(5) / 5, mode='same')  # 简单平滑

    # 动态阈值计算
    proj_mean = np.mean(smoothed_projection)
    proj_std = np.std(smoothed_projection)
    proj_max = np.max(smoothed_projection)
    proj_min = np.min(smoothed_projection)
    proj_range = proj_max - proj_min

    if proj_range > 50:
        threshold = proj_mean + 0.8 * proj_std
    elif proj_range > 30:
        threshold = proj_mean + 0.6 * proj_std
    else:
        threshold = proj_mean + 0.4 * proj_std

    # 峰值检测
    peaks, properties = find_peaks(
        smoothed_projection,
        height=threshold,
        distance=max(10, w / 20),  # 最小距离动态调整
        prominence=proj_std
    )

    # 结果计算与展示
    if len(peaks) >= 2:
        pixel_spacings = np.diff(peaks)
        avg_pixel_spacing = np.median(pixel_spacings)  # 使用中位数更鲁棒
        actual_spacing = avg_pixel_spacing * (scale_bar_actual / scale_bar_pixels)

        # 计算标准差和相对误差
        std_dev = np.std(pixel_spacings)
        rel_error = (std_dev / avg_pixel_spacing) * 100

        print(f"检测到 {len(peaks)} 个条纹")
        print(f"平均像素间隔: {avg_pixel_spacing:.2f} ± {std_dev:.2f} 像素 (相对误差: {rel_error:.1f}%)")
        print(f"实际物理间隔: {actual_spacing:.2f} nm")

        # 可视化结果
        plt.figure(figsize=(12, 8))
        plt.subplot(211), plt.imshow(rotated_roi, cmap='gray')
        plt.title(f'旋转后ROI - 间隔:{actual_spacing:.2f} nm')

        plt.subplot(212)
        plt.plot(smoothed_projection, label='平滑投影')
        plt.plot(projection, alpha=0.3, label='原始投影')
        plt.axhline(y=threshold, color='r', linestyle='--', label='阈值')
        plt.plot(peaks, smoothed_projection[peaks], 'ro', label='条纹峰值')

        # 标记间距
        for i in range(len(peaks) - 1):
            mid_point = (peaks[i] + peaks[i + 1]) // 2
            plt.annotate(f'{pixel_spacings[i]:.1f}',
                         xy=(mid_point, smoothed_projection[mid_point]),
                         xytext=(0, 10),
                         textcoords='offset points',
                         ha='center')

        plt.title(f'{projection_dir}投影结果')
        plt.xlabel('位置')
        plt.ylabel('投影强度')
        plt.legend()
        plt.grid(True)
        plt.tight_layout()
        plt.show()

        if rel_error > 15:
            print("警告:测量误差较大，建议检查ROI选择或手动测量")
        elif len(peaks) < 5:
            print("提示:检测到的条纹数量较少，结果可能不够精确")
    else:
        print("错误:未检测到足够条纹，请尝试以下方法:")
        print("1. 重新选择ROI区域")
        print("2. 调整旋转角度")
        print("3. 尝试手动模式")

        # 可视化投影曲线帮助诊断
        plt.figure(figsize=(10, 4))
        plt.plot(smoothed_projection)
        plt.axhline(y=threshold, color='r', linestyle='--', label='阈值')
        plt.title(f'{projection_dir}投影曲线 (未检测到足够峰值)')
        plt.xlabel('位置')
        plt.ylabel('投影强度')
        plt.legend()
        plt.grid(True)
        plt.show()


# 运行示例
if __name__ == "__main__":
    image_path = "../img/pic3.jpg"  # 图像路径
    scale_bar_pixels = 50  # 比例尺像素长度
    scale_bar_actual = 10  # 比例尺实际长度(nm)
    manual_mode = False  # 是否手动输入角度

    # 设置中文字体支持
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

    measure_stripe_spacing(image_path, scale_bar_pixels, scale_bar_actual, manual_mode)